package syncx

import (
	"sync"
	"sync/atomic"
)

type singleLock struct {
	counter int64
	lock    *sync.RWMutex
}

// MultipleLock is the main interface for lock base on key
type MultipleLock interface {
	// Lock base on the key
	Lock(interface{})

	// RLock lock the rw for reading
	RLock(interface{})

	// Unlock the key
	Unlock(interface{})

	// RUnlock the the read lock
	RUnlock(interface{})
}

// A multi lock type
type lock struct {
	// Why is a lock needed here?
	// Consider a scenario where A is about to release the lock, B and C have acquired the lock before releasing, but there is no incr counter yet.
	// At this time, A decr counter thinks that count=0 and deletes this key from inUse. Suppose it is locked in B's hand at this time, and B is ready to go
	// Release the lock. The locker taken out at this time is no longer the previous locker, and the previous lock has not been released (resulting in subsequent
	// key gets a locked lock from the pool), C will always block on Lock. The meaning of this mu is to get
	// signalLock and counter counting become atomic operations.
	mu sync.Mutex

	inUse sync.Map
	pool  *sync.Pool
}

func (l *lock) Lock(key interface{}) {
	l.mu.Lock()
	m := l.getLocker(key)
	atomic.AddInt64(&m.counter, 1)
	l.mu.Unlock()
	m.lock.Lock()
}

func (l *lock) RLock(key interface{}) {
	l.mu.Lock()
	m := l.getLocker(key)
	atomic.AddInt64(&m.counter, 1)
	l.mu.Unlock()
	m.lock.RLock()
}

func (l *lock) Unlock(key interface{}) {
	l.mu.Lock()
	m := l.getLocker(key)
	m.lock.Unlock()
	l.putBackInPool(key, m)
	l.mu.Unlock()
}

func (l *lock) RUnlock(key interface{}) {
	l.mu.Lock()
	m := l.getLocker(key)
	m.lock.RUnlock()
	l.putBackInPool(key, m)
	l.mu.Unlock()
}

func (l *lock) putBackInPool(key interface{}, m *singleLock) {
	atomic.AddInt64(&m.counter, -1)
	if m.counter <= 0 {
		l.pool.Put(m.lock)
		l.inUse.Delete(key)
	}
}

func (l *lock) getLocker(key interface{}) *singleLock {
	res, _ := l.inUse.LoadOrStore(key, &singleLock{
		counter: 0,
		lock:    l.pool.Get().(*sync.RWMutex),
	})

	return res.(*singleLock)
}

// NewMultipleLock create a new multiple lock
func NewMultipleLock() MultipleLock {
	return &lock{
		pool: &sync.Pool{
			New: func() interface{} {
				return &sync.RWMutex{}
			},
		},
	}
}
